import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Any, Optional
import datetime
from model_persistence import ModelPersistenceManager

class GoldTradingGeneticAlgorithm:
    """黄金交易遗传算法模型"""
    
    def __init__(self, population_size: int = 100, generations: int = 200):
        self.population_size = population_size
        self.generations = generations
        self.population = []
        self.best_individual = None
        self.best_fitness = float('-inf')
        self.fitness_history = []
        self.training_start_time = None
        self.training_end_time = None
        
        # 黄金交易特定参数
        self.gold_specific_params = {
            'volatility_weight': 0.3,
            'trend_weight': 0.4,
            'support_resistance_weight': 0.3,
            'risk_reward_ratio': 2.0,
            'max_drawdown_threshold': 0.05
        }
    
    def initialize_population(self):
        """初始化种群"""
        self.population = []
        for _ in range(self.population_size):
            individual = {
                'ma_short': np.random.randint(5, 50),
                'ma_long': np.random.randint(20, 200),
                'rsi_period': np.random.randint(10, 30),
                'rsi_oversold': np.random.uniform(20, 35),
                'rsi_overbought': np.random.uniform(65, 80),
                'stop_loss_pct': np.random.uniform(0.01, 0.05),
                'take_profit_pct': np.random.uniform(0.02, 0.10),
                'position_size_pct': np.random.uniform(0.01, 0.05)
            }
            self.population.append(individual)
    
    def fitness_function(self, individual: Dict, historical_data: pd.DataFrame) -> float:
        """
        适应度函数 - 评估个体在黄金交易中的表现
        
        Args:
            individual: 个体参数
            historical_data: 历史数据
        
        Returns:
            适应度分数
        """
        # 模拟交易回测
        returns = self.backtest_strategy(individual, historical_data)
        
        if len(returns) == 0:
            return float('-inf')
        
        # 计算夏普比率
        sharpe_ratio = np.mean(returns) / (np.std(returns) + 1e-8) * np.sqrt(252)
        
        # 计算最大回撤
        cumulative_returns = np.cumprod(1 + returns)
        running_max = np.maximum.accumulate(cumulative_returns)
        drawdown = (cumulative_returns - running_max) / running_max
        max_drawdown = np.min(drawdown)
        
        # 计算胜率
        win_rate = np.sum(returns > 0) / len(returns)
        
        # 综合适应度分数
        fitness = (
            sharpe_ratio * 0.4 +
            (1 + max_drawdown) * 0.3 +  # 最大回撤越小越好
            win_rate * 0.3
        )
        
        return fitness
    
    def backtest_strategy(self, individual: Dict, data: pd.DataFrame) -> List[float]:
        """简单的策略回测"""
        returns = []
        position = 0
        entry_price = 0
        
        for i in range(len(data)):
            # 简化的交易逻辑
            if position == 0:  # 无持仓
                if data.iloc[i]['close'] > data.iloc[i]['ma_short']:
                    position = 1
                    entry_price = data.iloc[i]['close']
            elif position == 1:  # 持多仓
                # 止损或止盈
                current_return = (data.iloc[i]['close'] - entry_price) / entry_price
                if (current_return <= -individual['stop_loss_pct'] or 
                    current_return >= individual['take_profit_pct']):
                    returns.append(current_return)
                    position = 0
        
        return returns
    
    def train(self, historical_data: pd.DataFrame) -> Dict[str, Any]:
        """
        训练遗传算法模型
        
        Args:
            historical_data: 历史交易数据
        
        Returns:
            训练结果和元数据
        """
        self.training_start_time = datetime.datetime.now()
        
        # 初始化种群
        self.initialize_population()
        
        # 训练循环
        for generation in range(self.generations):
            # 评估适应度
            fitness_scores = []
            for individual in self.population:
                fitness = self.fitness_function(individual, historical_data)
                fitness_scores.append(fitness)
            
            # 更新最佳个体
            best_idx = np.argmax(fitness_scores)
            if fitness_scores[best_idx] > self.best_fitness:
                self.best_fitness = fitness_scores[best_idx]
                self.best_individual = self.population[best_idx].copy()
            
            self.fitness_history.append(self.best_fitness)
            
            # 选择、交叉、变异
            self.selection(fitness_scores)
            self.crossover()
            self.mutation()
        
        self.training_end_time = datetime.datetime.now()
        
        # 准备训练元数据
        metadata = {
            'best_fitness': self.best_fitness,
            'best_individual': self.best_individual,
            'fitness_history': self.fitness_history,
            'training_time': str(self.training_end_time - self.training_start_time),
            'population_size': self.population_size,
            'generations': self.generations,
            'gold_specific_params': self.gold_specific_params
        }
        
        return metadata
    
    def selection(self, fitness_scores: List[float]):
        """选择操作"""
        # 轮盘赌选择
        fitness_array = np.array(fitness_scores)
        fitness_array = fitness_array - np.min(fitness_array) + 1e-8
        probabilities = fitness_array / np.sum(fitness_array)
        
        selected_indices = np.random.choice(
            len(self.population), 
            size=len(self.population), 
            p=probabilities
        )
        
        self.population = [self.population[i] for i in selected_indices]
    
    def crossover(self):
        """交叉操作"""
        for i in range(0, len(self.population), 2):
            if i + 1 < len(self.population):
                parent1 = self.population[i]
                parent2 = self.population[i + 1]
                
                # 单点交叉
                crossover_point = len(parent1) // 2
                keys = list(parent1.keys())
                
                child1 = {}
                child2 = {}
                
                for j, key in enumerate(keys):
                    if j < crossover_point:
                        child1[key] = parent1[key]
                        child2[key] = parent2[key]
                    else:
                        child1[key] = parent2[key]
                        child2[key] = parent1[key]
                
                self.population[i] = child1
                self.population[i + 1] = child2
    
    def mutation(self):
        """变异操作"""
        mutation_rate = 0.1
        
        for individual in self.population:
            if np.random.random() < mutation_rate:
                # 随机选择一个参数进行变异
                key = np.random.choice(list(individual.keys()))
                
                if 'ma_' in key:
                    individual[key] = np.random.randint(5, 200)
                elif 'rsi_period' in key:
                    individual[key] = np.random.randint(10, 30)
                elif 'rsi_' in key:
                    individual[key] = np.random.uniform(20, 80)
                else:
                    individual[key] = np.random.uniform(0.01, 0.10)

class GoldTradingRuleForest:
    """黄金交易规则森林模型"""
    
    def __init__(self, n_rules: int = 50, min_support: float = 0.1):
        self.n_rules = n_rules
        self.min_support = min_support
        self.rules = []
        self.rule_weights = []
        self.accuracy = 0.0
        self.training_start_time = None
        self.training_end_time = None
        
        # 黄金交易特定规则模板
        self.gold_rule_templates = [
            "IF price > ma_short AND rsi < 70 THEN buy",
            "IF price < ma_long AND rsi > 30 THEN sell",
            "IF volume > avg_volume * 1.5 AND price_up THEN buy",
            "IF support_level < price < resistance_level THEN hold",
            "IF volatility > threshold THEN reduce_position"
        ]
    
    def generate_rules(self, data: pd.DataFrame) -> List[Dict]:
        """生成交易规则"""
        rules = []
        
        # 技术指标规则
        for _ in range(self.n_rules // 2):
            rule = self.generate_technical_rule(data)
            if rule:
                rules.append(rule)
        
        # 价格模式规则
        for _ in range(self.n_rules // 2):
            rule = self.generate_pattern_rule(data)
            if rule:
                rules.append(rule)
        
        return rules
    
    def generate_technical_rule(self, data: pd.DataFrame) -> Optional[Dict]:
        """生成技术指标规则"""
        rule_types = ['ma_cross', 'rsi_signal', 'volume_signal', 'volatility_signal']
        rule_type = np.random.choice(rule_types)
        
        if rule_type == 'ma_cross':
            ma_short = np.random.randint(5, 20)
            ma_long = np.random.randint(20, 50)
            return {
                'type': 'ma_cross',
                'condition': f'ma_{ma_short} > ma_{ma_long}',
                'action': 'buy',
                'confidence': np.random.uniform(0.6, 0.9)
            }
        
        elif rule_type == 'rsi_signal':
            rsi_period = np.random.randint(10, 30)
            rsi_threshold = np.random.uniform(30, 70)
            return {
                'type': 'rsi_signal',
                'condition': f'rsi_{rsi_period} < {rsi_threshold}',
                'action': 'buy',
                'confidence': np.random.uniform(0.6, 0.9)
            }
        
        return None
    
    def generate_pattern_rule(self, data: pd.DataFrame) -> Optional[Dict]:
        """生成价格模式规则"""
        patterns = ['support_break', 'resistance_break', 'trend_reversal', 'consolidation']
        pattern = np.random.choice(patterns)
        
        if pattern == 'support_break':
            return {
                'type': 'support_break',
                'condition': 'price < support_level * 0.98',
                'action': 'sell',
                'confidence': np.random.uniform(0.7, 0.95)
            }
        
        elif pattern == 'resistance_break':
            return {
                'type': 'resistance_break',
                'condition': 'price > resistance_level * 1.02',
                'action': 'buy',
                'confidence': np.random.uniform(0.7, 0.95)
            }
        
        return None
    
    def evaluate_rule(self, rule: Dict, data: pd.DataFrame) -> float:
        """评估规则的有效性"""
        # 简化的规则评估
        correct_predictions = 0
        total_predictions = 0
        
        for i in range(len(data) - 1):
            # 检查规则条件
            if self.check_rule_condition(rule, data.iloc[i]):
                total_predictions += 1
                
                # 检查预测是否正确
                if self.check_prediction_accuracy(rule, data.iloc[i], data.iloc[i + 1]):
                    correct_predictions += 1
        
        if total_predictions == 0:
            return 0.0
        
        return correct_predictions / total_predictions
    
    def check_rule_condition(self, rule: Dict, data_row: pd.Series) -> bool:
        """检查规则条件是否满足"""
        # 简化的条件检查
        condition = rule['condition']
        
        if 'ma_' in condition:
            return True  # 简化处理
        elif 'rsi_' in condition:
            return True  # 简化处理
        elif 'support_level' in condition:
            return True  # 简化处理
        
        return False
    
    def check_prediction_accuracy(self, rule: Dict, current_data: pd.Series, next_data: pd.Series) -> bool:
        """检查预测准确性"""
        action = rule['action']
        
        if action == 'buy':
            return next_data['close'] > current_data['close']
        elif action == 'sell':
            return next_data['close'] < current_data['close']
        
        return False
    
    def train(self, historical_data: pd.DataFrame) -> Dict[str, Any]:
        """
        训练规则森林模型
        
        Args:
            historical_data: 历史交易数据
        
        Returns:
            训练结果和元数据
        """
        self.training_start_time = datetime.datetime.now()
        
        # 生成规则
        self.rules = self.generate_rules(historical_data)
        
        # 评估规则
        rule_scores = []
        for rule in self.rules:
            score = self.evaluate_rule(rule, historical_data)
            rule_scores.append(score)
        
        # 选择最佳规则
        best_rules = []
        best_scores = []
        
        for i, score in enumerate(rule_scores):
            if score > self.min_support:
                best_rules.append(self.rules[i])
                best_scores.append(score)
        
        self.rules = best_rules
        self.rule_weights = best_scores
        
        # 计算整体准确率
        if len(self.rules) > 0:
            self.accuracy = np.mean(self.rule_weights)
        
        self.training_end_time = datetime.datetime.now()
        
        # 准备训练元数据
        metadata = {
            'n_rules': len(self.rules),
            'accuracy': self.accuracy,
            'rule_weights': self.rule_weights,
            'training_time': str(self.training_end_time - self.training_start_time),
            'min_support': self.min_support,
            'best_rules': self.rules[:10]  # 保存前10个最佳规则
        }
        
        return metadata
    
    def predict(self, current_data: pd.DataFrame) -> Dict[str, float]:
        """
        使用规则森林进行预测
        
        Args:
            current_data: 当前市场数据
        
        Returns:
            预测结果
        """
        predictions = {'buy': 0.0, 'sell': 0.0, 'hold': 0.0}
        
        for i, rule in enumerate(self.rules):
            if self.check_rule_condition(rule, current_data.iloc[0]):
                action = rule['action']
                weight = self.rule_weights[i]
                predictions[action] += weight
        
        # 归一化
        total_weight = sum(predictions.values())
        if total_weight > 0:
            for action in predictions:
                predictions[action] /= total_weight
        
        return predictions

class GoldTradingModelManager:
    """黄金交易模型管理器"""
    
    def __init__(self, models_path: str = "./gold_models"):
        self.persistence_manager = ModelPersistenceManager(models_path)
        self.genetic_model = None
        self.rule_forest_model = None
    
    def train_and_save_models(self, historical_data: pd.DataFrame, 
                            model_name: str = "gold_trading") -> Dict[str, str]:
        """
        训练并保存模型
        
        Args:
            historical_data: 历史数据
            model_name: 模型名称
        
        Returns:
            保存的文件路径
        """
        saved_paths = {}
        
        # 训练遗传算法模型
        print("开始训练遗传算法模型...")
        self.genetic_model = GoldTradingGeneticAlgorithm()
        ga_metadata = self.genetic_model.train(historical_data)
        ga_path = self.persistence_manager.save_genetic_model(
            self.genetic_model, f"{model_name}_ga", ga_metadata
        )
        saved_paths['genetic_algorithm'] = ga_path
        
        # 训练规则森林模型
        print("开始训练规则森林模型...")
        self.rule_forest_model = GoldTradingRuleForest()
        rf_metadata = self.rule_forest_model.train(historical_data)
        rf_path = self.persistence_manager.save_rule_forest_model(
            self.rule_forest_model, f"{model_name}_rf", rf_metadata
        )
        saved_paths['rule_forest'] = rf_path
        
        return saved_paths
    
    def load_models(self, model_name: str = "gold_trading") -> bool:
        """
        加载已训练的模型
        
        Args:
            model_name: 模型名称
        
        Returns:
            是否成功加载
        """
        # 加载遗传算法模型
        ga_data = self.persistence_manager.load_latest_model(
            'genetic_algorithm', f"{model_name}_ga"
        )
        if ga_data:
            self.genetic_model = ga_data['model']
            print(f"遗传算法模型加载成功，最佳适应度: {ga_data['metadata']['best_fitness']:.4f}")
        else:
            print("遗传算法模型加载失败")
            return False
        
        # 加载规则森林模型
        rf_data = self.persistence_manager.load_latest_model(
            'rule_forest', f"{model_name}_rf"
        )
        if rf_data:
            self.rule_forest_model = rf_data['model']
            print(f"规则森林模型加载成功，准确率: {rf_data['metadata']['accuracy']:.4f}")
        else:
            print("规则森林模型加载失败")
            return False
        
        return True
    
    def ensemble_predict(self, current_data: pd.DataFrame) -> Dict[str, float]:
        """
        集成预测 - 结合遗传算法和规则森林的结果
        
        Args:
            current_data: 当前市场数据
        
        Returns:
            集成预测结果
        """
        if not self.genetic_model or not self.rule_forest_model:
            raise ValueError("模型未加载，请先调用load_models()")
        
        # 遗传算法预测
        ga_prediction = self.genetic_model.best_individual
        
        # 规则森林预测
        rf_prediction = self.rule_forest_model.predict(current_data)
        
        # 集成预测（简单平均）
        ensemble_prediction = {
            'buy': rf_prediction['buy'] * 0.6,  # 规则森林权重60%
            'sell': rf_prediction['sell'] * 0.6,
            'hold': rf_prediction['hold'] * 0.6
        }
        
        # 根据遗传算法参数调整
        if ga_prediction:
            # 这里可以根据遗传算法的最佳参数进行进一步调整
            pass
        
        return ensemble_prediction

# 使用示例
if __name__ == "__main__":
    # 创建模型管理器
    manager = GoldTradingModelManager()
    
    # 模拟历史数据
    dates = pd.date_range('2023-01-01', '2024-01-01', freq='D')
    np.random.seed(42)
    
    historical_data = pd.DataFrame({
        'date': dates,
        'open': np.random.normal(1800, 50, len(dates)),
        'high': np.random.normal(1820, 50, len(dates)),
        'low': np.random.normal(1780, 50, len(dates)),
        'close': np.random.normal(1800, 50, len(dates)),
        'volume': np.random.normal(1000, 200, len(dates))
    })
    
    # 添加技术指标
    historical_data['ma_short'] = historical_data['close'].rolling(10).mean()
    historical_data['ma_long'] = historical_data['close'].rolling(50).mean()
    historical_data['rsi'] = 50 + np.random.normal(0, 10, len(dates))
    
    # 训练并保存模型
    print("开始训练模型...")
    saved_paths = manager.train_and_save_models(historical_data, "gold_trading_v1")
    print("模型训练完成，保存路径:", saved_paths)
    
    # 加载模型
    print("\n开始加载模型...")
    success = manager.load_models("gold_trading_v1")
    
    if success:
        # 进行预测
        current_data = historical_data.tail(1)
        prediction = manager.ensemble_predict(current_data)
        print(f"集成预测结果: {prediction}")
    
    # 列出所有可用模型
    available_models = manager.persistence_manager.list_available_models()
    print(f"\n可用模型数量: {len(available_models.get('genetic_algorithm', []))} 个遗传算法模型")
    print(f"可用模型数量: {len(available_models.get('rule_forest', []))} 个规则森林模型") 